#include "mpi_util.hpp"
#include <iostream>

void MPI_Send_String(const std::string &buf, int dest, int tag, MPI_Comm comm) {
  MPI_Send(const_cast<char *>(buf.data()), buf.size(), MPI_CHAR, dest, tag,
           comm);
}
void MPI_Send_Strings(const std::vector<std::string> &buffs, int dest, int tag,
                      MPI_Comm comm) {
  for (auto i = 0; i < buffs.size(); i++) {
    MPI_Send_String(buffs[i], dest, tag, comm);
  }
}

void MPI_Recv_String(std::string &buf, int src, int tag, MPI_Comm comm) {
  MPI_Status status;
  MPI_Probe(src, tag, comm, &status);
  int size;
  MPI_Get_count(&status, MPI_CHAR, &size);
  buf.resize(size);
  MPI_Recv(const_cast<char *>(buf.data()), size, MPI_CHAR, src, tag, comm,
           MPI_STATUS_IGNORE);
}
void MPI_Recv_Strings(std::vector<std::string> &buffs, int vec_size, int src,
                      int tag, MPI_Comm comm) {
  buffs.resize(vec_size);
  for (int i = 0; i < buffs.size(); i++) {
    MPI_Status status;
    MPI_Probe(src, tag, comm, &status);
    int size;
    MPI_Get_count(&status, MPI_CHAR, &size);
    buffs[i].resize(size);
    MPI_Recv(const_cast<char *>(buffs[i].data()), size, MPI_CHAR, src, tag,
             comm, MPI_STATUS_IGNORE);
  }
}
void MPI_PackRecv_Strings(std::vector<std::string> &buffs, int vec_size,
                          int src, int tag, MPI_Comm comm) {
  int size_buff[vec_size];
  MPI_Recv(size_buff, vec_size, MPI_INT, src, tag, comm, MPI_STATUS_IGNORE);

  std::string pack_buff;
  MPI_Recv_String(pack_buff, src, tag, comm);

  // Unpack
  int start = 0;
  buffs.clear();
  for (int i = 0; i < vec_size; i++) {
    buffs.emplace_back(pack_buff.substr(start, size_buff[i]));
    start += size_buff[i];
  }
}

void MPI_MultiPackRecv_Strings(std::vector<std::string> &buffs, int vec_size,
                               int src, int tag, MPI_Comm comm) {
  buffs.clear();
  int unpack_size = 0;
  while (unpack_size < vec_size) {
    MPI_Status status;
    MPI_Probe(src, tag, comm, &status);
    int size;
    MPI_Get_count(&status, MPI_INT, &size);
    int size_buff[size];
    MPI_Recv(size_buff, size, MPI_INT, src, tag, comm, MPI_STATUS_IGNORE);

    std::string pack_buff;
    MPI_Recv_String(pack_buff, src, tag, comm);

    // Unpack
    int start = 0;
    for (int i = 0; i < size; i++) {
      buffs.emplace_back(pack_buff.substr(start, size_buff[i]));
      start += size_buff[i];
      unpack_size++;
    }
  }
}

void MPI_Bcast_String(std::string &buf, int root, MPI_Comm comm) {
  int my_rank;
  MPI_Comm_rank(comm, &my_rank);
  int size = buf.size();
  MPI_Bcast(&size, 1, MPI_INT, root, comm);
  if (my_rank != root) {
    buf.resize(size);
  }
  MPI_Bcast(const_cast<char *>(buf.data()), size, MPI_CHAR, root, comm);
}
void MPI_Bcast_Strings(std::vector<std::string> &buffs, int root,
                       MPI_Comm comm) {
  int my_rank;
  MPI_Comm_rank(comm, &my_rank);
  int vec_size = buffs.size();
  MPI_Bcast(&vec_size, 1, MPI_INT, root, comm);
  if (my_rank != root) {
    buffs.resize(vec_size);
  }
  for (auto i = 0; i < buffs.size(); i++) {
    int size = buffs[i].size();
    MPI_Bcast(&size, 1, MPI_INT, root, comm);
    if (my_rank != root) {
      buffs[i].resize(size);
    }
    MPI_Bcast(const_cast<char *>(buffs[i].data()), size, MPI_CHAR, root, comm);
  }
}

//! Assume that pack_buff.size() < INT_MAX
void MPI_PackBcast_Strings(std::vector<std::string> &buffs, int root,
                           MPI_Comm comm) {
  int my_rank;
  MPI_Comm_rank(comm, &my_rank);
  int vec_size = buffs.size();
  MPI_Bcast(&vec_size, 1, MPI_INT, root, comm);

  std::string pack_buff;
  std::vector<int> size_buff;
  if (my_rank != root) {
    buffs.resize(vec_size);
  }
  size_buff.resize(vec_size);
  if (my_rank == root) {
    // Pack
    for (int i = 0; i < vec_size; i++) {
      pack_buff += buffs[i];
      size_buff[i] = buffs[i].size();
    }
  }

  // Bcast: size_buff
  MPI_Bcast(size_buff.data(), size_buff.size(), MPI_INT, root, comm);

  // Bcast: pack_buff
  int size = pack_buff.size();
  MPI_Bcast(&size, 1, MPI_INT, root, comm);
  if (my_rank != root) {
    pack_buff.resize(size);
  }
  MPI_Bcast(const_cast<char *>(pack_buff.data()), size, MPI_CHAR, root, comm);

  // Unpack
  int start = 0;
  buffs.clear();
  for (int i = 0; i < vec_size; i++) {
    buffs.emplace_back(pack_buff.substr(start, size_buff[i]));
    start += size_buff[i];
  }
}

void MPI_PackIbcast_Strings(std::vector<std::string> &buffs, int root,
                            MPI_Comm comm) {
  if (comm == MPI_COMM_NULL) {
    return;
  }
  int my_rank, comm_size;
  MPI_Comm_rank(comm, &my_rank);
  MPI_Comm_size(comm, &comm_size);
  if (comm_size == 1) {
    return;
  }

  std::vector<std::string> pack_buffs;
  std::vector<std::vector<int>> size_buffs;
  int pack_id = 0;
  if (my_rank == root) {
    int i = 0;
    while (i < buffs.size()) {
      std::vector<int> size_buff;
      std::string pack_buff;
      int cur_size = 0;
      int j = i;
      while (j < buffs.size() &&
             cur_size + static_cast<int>(buffs[j].size()) > 0) {
        size_buff.emplace_back(buffs[j].size());
        pack_buff += buffs[j];
        cur_size += static_cast<int>(buffs[j].size());
        j++;
      }
      i = j;
      size_buffs.emplace_back(size_buff);
      pack_buffs.emplace_back(pack_buff);
      pack_id++;
    }
  }

  MPI_Bcast(&pack_id, 1, MPI_INT, root, comm);
  if (my_rank != root) {
    size_buffs.resize(pack_id);
    pack_buffs.resize(pack_id);
  }
  std::vector<MPI_Request> reqs(pack_id, MPI_REQUEST_NULL);
  for (int i = 0; i < pack_id; i++) {
    int sz1, sz2;
    if (my_rank == root) {
      sz1 = size_buffs[i].size();
      sz2 = pack_buffs[i].size();
    }
    MPI_Bcast(&sz1, 1, MPI_INT, root, comm);
    MPI_Bcast(&sz2, 1, MPI_INT, root, comm);
    if (my_rank != root) {
      size_buffs[i].resize(sz1);
      pack_buffs[i].resize(sz2);
    }
    MPI_Ibcast(size_buffs[i].data(), size_buffs[i].size(), MPI_INT, root, comm,
               &reqs[i]);
    MPI_Ibcast(const_cast<char *>(pack_buffs[i].data()), pack_buffs[i].size(),
               MPI_CHAR, root, comm, &reqs[i]);
  }

  MPI_Waitall(pack_id, reqs.data(), MPI_STATUSES_IGNORE);

  // Unpack
  if (my_rank != root) {
    buffs.clear();
    for (int k = 0; k < pack_id; k++) {
      int start = 0;
      auto &size_buff = size_buffs[k];
      auto &pack_buff = pack_buffs[k];
      for (int i = 0; i < size_buff.size(); i++) {
        buffs.emplace_back(pack_buff.substr(start, size_buff[i]));
        start += size_buff[i];
      }
    }
  }
}

// Safe
void MPI_Isendrecv_String_replace(std::string &buf, int dest, int tag_send,
                                  int src, int tag_recv, MPI_Comm comm) {
  int size = buf.size();
  MPI_Sendrecv_replace(&size, 1, MPI_INT, dest, tag_send, src, tag_recv, comm,
                       MPI_STATUS_IGNORE);
  std::string recv;
  recv.resize(size);
  MPI_Request reqs[2];
  MPI_Isend(const_cast<char *>(buf.data()), buf.size(), MPI_CHAR, dest,
            tag_send, comm, &reqs[0]);
  MPI_Irecv(const_cast<char *>(recv.data()), size, MPI_CHAR, src, tag_recv,
            comm, &reqs[1]);
  MPI_Waitall(2, reqs, MPI_STATUSES_IGNORE);
  buf = std::move(recv);
}
//! Assume that buf.size() is the same
void MPI_Sendrecv_String_replace(std::string &buf, int dest, int tag_send,
                                 int src, int tag_recv, MPI_Comm comm) {
  MPI_Sendrecv_replace(const_cast<char *>(buf.data()), buf.size(), MPI_CHAR,
                       dest, tag_send, src, tag_recv, comm, MPI_STATUS_IGNORE);
}
void MPI_Sendrecv_String_replace(std::vector<std::string> &buffs, int dest,
                                 int tag_send, int src, int tag_recv,
                                 MPI_Comm comm) {
  for (auto i = 0; i < buffs.size(); i++) {
    MPI_Sendrecv_String_replace(buffs[i], dest, tag_send, src, tag_recv, comm);
  }
}

//* Non-blocking
void MPI_Isend_String(const std::string &buf, int dest, int tag, MPI_Comm comm,
                      MPI_Request *req) {
  MPI_Isend(const_cast<char *>(buf.data()), buf.size(), MPI_CHAR, dest, tag,
            comm, req);
}

void MPI_Isend_Strings(const std::vector<std::string> &buffs, int dest, int tag,
                       MPI_Comm comm, MPI_Request *req) {
  for (auto i = 0; i < buffs.size(); i++) {
    MPI_Isend(const_cast<char *>(buffs[i].data()), buffs[i].size(), MPI_CHAR,
              dest, tag, comm, &req[i]);
  }
}
void MPI_PackIsend_Strings(const std::vector<std::string> &buffs,
                           std::string &pack_buff, int *size_buff, int dest,
                           int tag, MPI_Comm comm, MPI_Request *req) {
  int cur_size = 0;
  for (auto i = 0; i < buffs.size(); i++) {
    if (cur_size + static_cast<int>(buffs[i].size()) < 0) {
      std::cerr << "Packing length overflow: cur_size + buffs[i].size = "
                << cur_size + static_cast<int>(buffs[i].size())
                << " < 0, cur_size = " << cur_size << ", i = " << i
                << std::endl;
      MPI_Abort(comm, 1);
    }
    cur_size += buffs[i].size();
    size_buff[i] = buffs[i].size();
    pack_buff += buffs[i];
  }
  MPI_Isend(size_buff, buffs.size(), MPI_INT, dest, tag, comm, req);
  MPI_Isend(const_cast<char *>(pack_buff.data()), pack_buff.size(), MPI_CHAR,
            dest, tag, comm, req);
}

void MPI_PackSend_Strings(const std::vector<std::string> &buffs,
                           std::string &pack_buff, int *size_buff, int dest,
                           int tag, MPI_Comm comm) {
  int cur_size = 0;
  for (auto i = 0; i < buffs.size(); i++) {
    if (cur_size + static_cast<int>(buffs[i].size()) < 0) {
      std::cerr << "Packing length overflow: cur_size + buffs[i].size = "
                << cur_size + static_cast<int>(buffs[i].size())
                << " < 0, cur_size = " << cur_size << ", i = " << i
                << std::endl;
      MPI_Abort(comm, 1);
    }
    cur_size += buffs[i].size();
    size_buff[i] = buffs[i].size();
    pack_buff += buffs[i];
  }
  MPI_Send(size_buff, buffs.size(), MPI_INT, dest, tag, comm);
  MPI_Send(const_cast<char *>(pack_buff.data()), pack_buff.size(), MPI_CHAR,
            dest, tag, comm);
}

void MPI_PackIsend_Strings(const std::vector<std::string> &buffs,
                           std::vector<std::string> &pack_buffs,
                           std::vector<std::vector<int>> &size_buffs, int dest,
                           int tag, MPI_Comm comm, MPI_Request *req) {
  int i = 0;
  int pack_id = 0;
  while (i < buffs.size()) {
    int cur_size = 0;
    std::string pack_buff;
    std::vector<int> size_buff;
    int j = i;
    while (j < buffs.size() &&
           cur_size + static_cast<int>(buffs[j].size()) > 0) {
      size_buff.emplace_back(buffs[j].size());
      pack_buff += buffs[j];
      cur_size += static_cast<int>(buffs[j].size());
      j++;
    }
    i = j;
    size_buffs.emplace_back(size_buff);
    pack_buffs.emplace_back(pack_buff);

    MPI_Isend(size_buffs[pack_id].data(), size_buffs[pack_id].size(), MPI_INT,
              dest, tag, comm, req);
    MPI_Isend(const_cast<char *>(pack_buffs[pack_id].data()),
              pack_buffs[pack_id].size(), MPI_CHAR, dest, tag, comm, req);
    pack_id++;
  }
}
